from PyQt5.QtWidgets import (
    QMessageBox, QWidget, QVBoxLayout, QHBoxLayout, QTabWidget,
    QFrame, QTextEdit, QCheckBox, QLineEdit, QComboBox, QPushButton, QFormLayout, QGroupBox
)
from PyQt5.QtCore import QTimer
import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure

from taskmanager import SerialManager


class PumpControl(QWidget):
    def __init__(self):
        super().__init__()
        self.pump_serial_manager = SerialManager(baudrate=115200)
        self.power_limit = 1000
        self.plot_timer = QTimer(self)
        self.initUI()
    
    def initUI(self):
        main_layout = QVBoxLayout()
        
        ## --------------- 1. 上侧: 设置 ---------------
        up_layout = QHBoxLayout()
        
        ## --------------- 1.1 第一列: 串口+泵设置 -----
        up1_layout = QVBoxLayout()
        # 串口选择 + 波特率配置
        com_group = QGroupBox("串口参数设置")
        com_layout = QFormLayout()
        self.serial_port_combo = QComboBox()
        self.serial_port_combo.addItems(self.pump_serial_manager.get_available_ports())
        com_layout.addRow("Serial Port", self.serial_port_combo)
        self.baudrate_combo = QComboBox()
        self.baudrate_combo.addItems(["115200", "9600"])
        com_layout.addRow("Baudrate", self.baudrate_combo)
        com_group.setLayout(com_layout)
        up1_layout.addWidget(com_group)
        
        # 串口按钮
        com_btn_layout = QHBoxLayout()
        self.refresh_button = QPushButton("Refresh")
        self.refresh_button.clicked.connect(self.on_refresh_button_clicked)
        self.connect_button = QPushButton("Connect")
        self.connect_button.clicked.connect(self.on_connect_button_clicked)
        com_btn_layout.addWidget(self.refresh_button)
        com_btn_layout.addWidget(self.connect_button)
        up1_layout.addLayout(com_btn_layout)
        
        # Pump设置
        pump_group = QGroupBox("Pump设置")
        pump_layout = QFormLayout()
        lineEdit_power_limit = QLineEdit()
        lineEdit_power_limit.setText(str(self.power_limit))
        lineEdit_power_limit.textChanged.connect(self.on_power_limit_changed)
        pump_layout.addRow("Power Limit(0-1400mW)", lineEdit_power_limit)
        
        pump_group.setLayout(pump_layout)
        up1_layout.addWidget(pump_group)
        
        # Pump按钮
        pump_btn_layout = QHBoxLayout()
        self.pump_start_button = QPushButton("Start Pump")
        self.pump_start_button.clicked.connect(self.on_pump_start_button_clicked)
        pump_btn_layout.addWidget(self.pump_start_button)
        
        up1_layout.addLayout(pump_btn_layout)
        
        ## --------------- 1.2 第二列: 控制方式 -----
        # 控制方式
        control_group = QGroupBox("控制方式")
        control_layout = QVBoxLayout()
        control_tab_widget = QTabWidget()

        # 3种控制方式
        tab_power_control = QWidget()
        tab_PID_control = QWidget()
        tab_bang_bang_control = QWidget()
        control_tab_widget.addTab(tab_power_control, "Power Control")
        control_tab_widget.addTab(tab_PID_control, "PID Control")
        control_tab_widget.addTab(tab_bang_bang_control, "Bang Bang Control")
        self.on_control_tab_changed(0)
        control_tab_widget.currentChanged.connect(self.on_control_tab_changed)
        
        # Manual Control
        power_control_layout = QFormLayout()
        self.power_control_manualModeb = QComboBox()
        self.power_control_manualModeb.addItems(["Set val (register 23)", "Analog 1† (dial on evak kit)", "Analog 2† (pr. sense on eval kit)", "Analog 3 (main analog input)"])
        self.power_control_manualModeb.currentIndexChanged.connect(self.on_power_control_mode_changed)
        power_control_layout.addRow("Manual Mode", self.power_control_manualModeb)
        
        self.power_control_lineEdit = QLineEdit()
        self.power_control_lineEdit.setText("250")
        self.power_control_lineEdit.textChanged.connect(self.on_power_control_changed)
        power_control_layout.addRow("Power (mW)", self.power_control_lineEdit)
        tab_power_control.setLayout(power_control_layout)
        
        
        control_layout.addWidget(control_tab_widget)
        control_group.setLayout(control_layout)
        
        
        ## --------------- 1.3 第三列: 画图设置 -----
        up3_layout = QVBoxLayout()
        plot_group = QGroupBox("画图设置")
        plot_layout = QFormLayout()
        
        self.power_checkbox = QCheckBox("功率 (Power)")
        self.power_checkbox.setChecked(True)
        self.power_checkbox.stateChanged.connect(self.on_plot_setting_changed)
        plot_layout.addRow(self.power_checkbox)
        
        self.voltage_checkbox = QCheckBox("电压 (Voltage)")
        self.voltage_checkbox.setChecked(True)  # 默认勾选
        self.voltage_checkbox.stateChanged.connect(self.on_plot_setting_changed)
        plot_layout.addRow(self.voltage_checkbox)

        self.current_checkbox = QCheckBox("电流 (Current)")
        self.current_checkbox.setChecked(True)  # 默认勾选
        self.current_checkbox.stateChanged.connect(self.on_plot_setting_changed)
        plot_layout.addRow(self.current_checkbox)
        
        self.plot_time_lineEdit = QLineEdit()
        self.plot_time_lineEdit.setText("30")
        plot_layout.addRow("Plot Time (s)", self.plot_time_lineEdit)
        
        self.start_plot_button = QPushButton("Start Plot")
        self.start_plot_button.clicked.connect(self.on_start_plot_button_clicked)
        plot_layout.addRow(self.start_plot_button)
        
        plot_group.setLayout(plot_layout)
        up3_layout.addWidget(plot_group)
        
        up_layout.addLayout(up1_layout)
        up_layout.addWidget(control_group)
        up_layout.addLayout(up3_layout)
        
        ## --------------- 2. 下侧: 曲线图 -------------
        down_layout = QVBoxLayout()
        
        self.figure = Figure()
        self.canvas = FigureCanvas(self.figure)
        self.canvas.draw()
        
        down_layout.addWidget(self.canvas)
        
        main_layout.addLayout(up_layout)
        main_layout.addLayout(down_layout)
        self.setLayout(main_layout)

    # 按键: Refresh
    def on_refresh_button_clicked(self):
        # print("Updating serial ports")
        self.serial_port_combo.clear()
        # print(SerialManager.get_available_ports())
        self.serial_port_combo.addItems(self.pump_serial_manager.get_available_ports())
        print("Serial ports refreshed.")
    
    # 按键: Connect
    def on_connect_button_clicked(self):
        selected_port = self.serial_port_combo.currentText()
        if self.pump_serial_manager.connected:
            self.pump_serial_manager.disconnect()
            self.connect_button.setText("Connect")
            print(f"Disconnected from {selected_port}.")
            print(f"Disconnected from {selected_port}.")
        else:
            print(f"Connecting to {selected_port}...")
            if self.pump_serial_manager.connect(self.serial_port_combo.currentText()):
                self.connect_button.setText("Disconnect")
                print(f"Connected to {selected_port}.")
                print(f"Successfully Connected!")
            else:
                print(f"Cannot connect to {selected_port}.")
                QMessageBox.warning(self, "Error", "Cannot connect to the serial port")
    
    # 按键: 使能/关闭泵
    def on_pump_start_button_clicked(self):
        if self.pump_start_button.text() == "Start Pump":
            self.pump_start_button.setText("Stop Pump")
            self.pump_serial_manager.pump_enable()
        else:
            self.pump_start_button.setText("Start Pump")
            self.pump_serial_manager.pump_disable()
            
    # 文本框: Power Limit
    def on_power_limit_changed(self, text):
        try:
            self.power_limit = int(text)
            self.pump_serial_manager.pump_set_powerLimit(self.power_limit) 
        except ValueError:
            self.power_limit = 1000
            self.sender().setText(str(self.power_limit))
        print(f"Power Limit changed to {self.power_limit}.")
    
    # Tab: 控制方式
    def on_control_tab_changed(self, index):
        if index == 0:
            self.pump_serial_manager.pump_set_mode("Manual")
        elif index == 1:
            self.pump_serial_manager.pump_set_mode("PID")
        elif index == 2:
            self.pump_serial_manager.pump_set_mode("Bang Bang")
        print(f"Control Tab changed to {index}.")
    
    # 下拉框: Manual Mode
    def on_power_control_mode_changed(self, index):
        self.pump_serial_manager.pump_set_manual_mode(index)
        
    # 文本框: Power Control
    def on_power_control_changed(self, text):
        try:
            power = int(text)
            if power < 0:
                power = 0
                self.sender().setText(str(power))
            elif power > self.power_limit:
                power = self.power_limit
                self.sender().setText(str(power))
            self.pump_serial_manager.pump_set_manual_power(power)
        except ValueError:
            power = 250
            self.sender().setText(str(power))
        print(f"Power Control changed to {power}.")
    
    # 勾选框: 画图设置
    def on_plot_setting_changed(self):
        print("Plot setting changed.")
    
    # 按键: Start Plot
    def on_start_plot_button_clicked(self):
        if self.start_plot_button.text() == "Start Plot":
            self.start_plot_button.setText("Stop Plot")
            self.plot_time_lineEdit.setEnabled(False)
            self.plot_time = int(self.plot_time_lineEdit.text())
            
            # 开启串口 streaming mode
            self.pump_serial_manager.pump_streaming_mode_enable()
            
            # 清空绘图区域
            self.figure.clear()
            ax = self.figure.add_subplot(111)

            # 绘制曲线
            num_points = int(self.plot_time * 60)
            self.time_data = [i / 60 for i in range(num_points)] # 60Hz
            self.power_data = [0]*num_points
            self.voltage_data = [0]*num_points
            self.current_data = [0]*num_points
            self.curve_power, = ax.plot(self.time_data, self.power_data, label="Power")
            self.curve_voltage, = ax.plot(self.time_data, self.voltage_data, label="Voltage")
            self.curve_current, = ax.plot(self.time_data, self.current_data, label="Current")
                        
            # 图表设置
            ax.legend()
            ax.set_xlim(0, self.plot_time)
            ax.set_xlabel("Time (s)")
            
            # 画图定时器
            self.plot_timer.timeout.connect(lambda: self.update_plot(ax))
            self.plot_timer.start(int(1000/60))
             
        else:
            self.start_plot_button.setText("Start Plot")
            self.plot_time_lineEdit.setEnabled(True)
            # 关闭串口 streaming mode
            self.pump_serial_manager.pump_streaming_mode_disable()
            self.plot_timer.stop()
    
    # 更新曲线图
    def update_plot(self, ax):
        fields = self.pump_serial_manager.pump_streaming_data()
        if fields is None:
            return
        
        voltage = float(fields[1])
        current = float(fields[2])
        power = voltage * current
        
        # 将新数据添加到绘图列表中
        self.power_data.append(power)
        self.voltage_data.append(voltage)
        self.current_data.append(current)

        # 如果数据点超出范围，移除最早的数据点
        if len(self.power_data) > len(self.time_data):
            self.power_data.pop(0)
            self.voltage_data.pop(0)
            self.current_data.pop(0)

        # 更新曲线数据
        self.curve_power.set_data(self.time_data, self.power_data)
        self.curve_voltage.set_data(self.time_data, self.voltage_data)
        self.curve_current.set_data(self.time_data, self.current_data)

        # 是否显示曲线
        if not self.power_checkbox.isChecked():
            self.curve_power.set_visible(False)
        else:
            self.curve_power.set_visible(True)
        if not self.voltage_checkbox.isChecked():
            self.curve_voltage.set_visible(False)
        else:
            self.curve_voltage.set_visible(True)
        if not self.current_checkbox.isChecked():
            self.curve_current.set_visible(False)
        else:
            self.curve_current.set_visible(True)
            
        # 调整y坐标轴范围
        valid_data = []
        if self.power_checkbox.isChecked() and self.power_data:
            valid_data.extend([x for x in self.power_data if x is not None and not np.isnan(x)])
        if self.voltage_checkbox.isChecked() and self.voltage_data:
            valid_data.extend([x for x in self.voltage_data if x is not None and not np.isnan(x)])
        if self.current_checkbox.isChecked() and self.current_data:
            valid_data.extend([x for x in self.current_data if x is not None and not np.isnan(x)])

        if valid_data:
            max_value = max(valid_data) * 1.1  # 留出 10% 的空间
            min_value = min(valid_data) * 0.9  # 留出 10% 的空间
            ax.set_ylim(min_value, max_value)
        else:
            ax.set_ylim(0, 1)
        
        # 刷新绘图
        self.canvas.draw()
        
        